# 1. 导入所需库
from sentence_transformers import SentenceTransformer
from sklearn.metrics.pairwise import cosine_similarity
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

# 2. 初始化中文BERT模型（自动下载并加载，首次运行需等待模型下载）
try:
    # 加载bert-base-chinese模型（适合中文文本的基础嵌入模型）
    model = SentenceTransformer('bert-base-chinese')
    print("模型加载成功！")
except Exception as e:
    print(f"模型加载失败：{str(e)}")
    print("建议检查网络连接，或尝试先通过 huggingface-cli login 登录")
    exit()

# 3. 定义待比较的句子（可根据需求替换为自己的句子）
# 示例1：单组句子对（比较2个句子的相似度）
sentence_pair = [
    "小猫在院子里追蝴蝶",  # 句子A
    "一只猫咪在庭院中追逐蝴蝶"  # 句子B
]

# 示例2：多组句子（比较1个参考句与多个候选句的相似度，如文本检索场景）
reference_sentence = "人工智能技术正在改变医疗行业"  # 参考句
candidate_sentences = [
    "AI技术正在革新医疗领域",
    "人工智能对教育行业影响深远",
    "医疗行业因人工智能而发生变革",
    "机器学习算法在金融领域的应用"
]


# 4. 核心函数：计算句子相似度（基于余弦相似度）
def calculate_sentence_similarity(sentences, model):
    """
    生成句子嵌入并计算相似度矩阵
    :param sentences: 待处理的句子列表（list）
    :param model: 已加载的SentenceTransformer模型
    :return: 句子嵌入矩阵（numpy数组）、相似度矩阵（numpy数组）
    """
    # 生成句子嵌入（将文本转为768维向量，bert-base-chinese默认输出维度）
    embeddings = model.encode(
        sentences,
        convert_to_tensor=False,  # 转为numpy数组（便于后续计算）
        show_progress_bar=False   # 不显示进度条（简化输出）
    )
    
    # 计算余弦相似度（取值范围：[-1,1]，越接近1表示相似度越高）
    similarity_matrix = cosine_similarity(embeddings)
    
    return embeddings, similarity_matrix


# 5. 功能1：单组句子对相似度计算
print("\n" + "="*50)
print("功能1：单组句子对相似度比较")
print("="*50)
# 计算单组句子对的相似度
pair_embeddings, pair_similarity = calculate_sentence_similarity(sentence_pair, model)
# 提取相似度值（相似度矩阵的[0,1]位置即句子A与句子B的相似度）
similarity_score = pair_similarity[0][1]
# 输出结果（保留4位小数，更直观）
print(f"句子A：{sentence_pair[0]}")
print(f"句子B：{sentence_pair[1]}")
print(f"相似度得分：{similarity_score:.4f}")
print(f"相似度等级：{'极高' if similarity_score >= 0.8 else '较高' if similarity_score >= 0.6 else '中等' if similarity_score >= 0.4 else '较低'}")


# 6. 功能2：多候选句与参考句的相似度排序（如文本匹配、检索场景）
print("\n" + "="*50)
print("功能2：多候选句与参考句相似度排序")
print("="*50)
# 构造包含参考句的完整句子列表
all_sentences = [reference_sentence] + candidate_sentences
# 计算相似度
multi_embeddings, multi_similarity = calculate_sentence_similarity(all_sentences, model)
# 提取参考句（索引0）与所有候选句的相似度
reference_similarities = multi_similarity[0][1:]  # 跳过参考句自身

# 整理结果为DataFrame（便于查看和排序）
result_df = pd.DataFrame({
    "候选句": candidate_sentences,
    "相似度得分": reference_similarities.round(4),
    "相似度等级": [
        '极高' if s >= 0.8 else '较高' if s >= 0.6 else '中等' if s >= 0.4 else '较低'
        for s in reference_similarities
    ]
})
# 按相似度降序排序（相似度最高的排在前面）
result_df_sorted = result_df.sort_values("相似度得分", ascending=False).reset_index(drop=True)
# 输出排序后的结果
print(f"参考句：{reference_sentence}")
print("\n候选句相似度排序：")
# print(result_df_sorted.to_string(index_label="排名"))
result_df_sorted.index.name = "排名"  # 给索引列命名为“排名”
print(result_df_sorted.to_string())   # 直接转字符串，会自动显示索引名

# 7. 功能3：相似度矩阵可视化（热力图，直观展示所有句子间的相似度）
print("\n" + "="*50)
print("功能3：相似度矩阵可视化（热力图）")
print("="*50)
# 设置中文字体（避免matplotlib中文乱码）
plt.rcParams['font.sans-serif'] = ['SimHei', 'DejaVu Sans']
plt.rcParams['axes.unicode_minus'] = False

# 创建热力图
fig, ax = plt.subplots(figsize=(10, 8))  # 设置图的大小
# 绘制热力图（cmap=YlOrRd：颜色从黄到红，越红相似度越高）
im = ax.imshow(multi_similarity, cmap='YlOrRd', aspect='auto', vmin=0, vmax=1)

# 设置坐标轴标签（用句子的前15个字符作为标签，避免过长）
tick_labels = [sent[:15] + "..." if len(sent) > 15 else sent for sent in all_sentences]
ax.set_xticks(range(len(all_sentences)))
ax.set_yticks(range(len(all_sentences)))
ax.set_xticklabels(tick_labels, rotation=45, ha='right')  # 横坐标标签旋转45度
ax.set_yticklabels(tick_labels)

# 在热力图每个格子中添加相似度数值（保留2位小数）
for i in range(len(all_sentences)):
    for j in range(len(all_sentences)):
        text = ax.text(j, i, f'{multi_similarity[i, j]:.2f}',
                       ha="center", va="center", color="black" if multi_similarity[i, j] < 0.7 else "white")

# 添加颜色条（解释颜色对应的相似度值）
cbar = ax.figure.colorbar(im, ax=ax)
cbar.ax.set_ylabel("余弦相似度", rotation=-90, va="bottom")

# 设置图表标题
ax.set_title("句子相似度矩阵热力图（bert-base-chinese）", fontsize=14, pad=20)
# 调整布局（避免标签被截断）
plt.tight_layout()
# 显示图表（运行代码后会弹出窗口）
plt.show()